import os
import numpy as np
from PIL import Image
import joblib
import gradio as gr

LABEL_MAP = {0: 'cat', 1: 'dog'}
LAZY_MODEL_PATH = os.path.join('models', 'lazy_best_model.joblib')


def preprocess_image_flat(pil_img: Image.Image) -> np.ndarray:
    img = pil_img.convert('L').resize((32, 32))
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return arr.reshape(1, -1)


def load_model():
    if os.path.exists(LAZY_MODEL_PATH):
        payload = joblib.load(LAZY_MODEL_PATH)
        return payload['model'], payload.get('best_name', 'Unknown')
    # 兼容旧文件名
    fallback = os.path.join('models', 'best_model.joblib')
    if os.path.exists(fallback):
        payload = joblib.load(fallback)
        return payload['model'], 'best_model(joblib)'
    raise FileNotFoundError('未找到模型，请先运行 lp_train.py 或 train_ensemble.py 进行训练。')


MODEL, BEST_NAME = load_model()


def predict(img: Image.Image):
    x = preprocess_image_flat(img)
    if hasattr(MODEL, 'predict_proba'):
        proba = MODEL.predict_proba(x)[0]
        out = {LABEL_MAP[i]: float(p) for i, p in enumerate(proba)}
    else:
        y = int(MODEL.predict(x)[0])
        out = {LABEL_MAP[0]: 0.0, LABEL_MAP[1]: 0.0}
        out[LABEL_MAP[y]] = 1.0
    return out


def app():
    title = f"Dog vs Cat - LazyPredict 最优模型: {BEST_NAME}"
    desc = "上传图片，使用 LazyPredict 选出的最佳模型进行 CPU 推理（flat 特征）。"

    iface = gr.Interface(
        fn=predict,
        inputs=gr.Image(type='pil', label='上传图片'),
        outputs=gr.Label(num_top_classes=2, label='预测概率'),
        title=title,
        description=desc,
        allow_flagging='never'
    )
    iface.launch()


if __name__ == '__main__':
    app()
